import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split

from mindspore.dataset import GeneratorDataset

np.random.seed(0)

def loadTrainData():
    cancer = load_breast_cancer()  # 加载乳腺癌数据
    X = cancer.data  # 加载乳腺癌判别特征
    y = cancer.target  # 两个TAG，y = 0时为阴性，y = 1时为阳性
    # 将数据集划分为训练集和测试集，测试集占比为0.2
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    X_train = X_train.astype(np.float32)
    X_test = X_test.astype(np.float32)
    return X_train, X_test, y_train, y_test

X_train, X_test, y_train, y_test = loadTrainData()

# numpy版本太高都不符合当前版本的mindspore的数据集加载要求，numpy==1.26.4（版本已经较新了）可以，超过2.0版本大可能会报错
trainDataset = GeneratorDataset(source=zip(X_train, y_train), column_names=['data', 'target'], shuffle=True)
testDataset = GeneratorDataset(source=zip(X_test, y_test), column_names=['data', 'target'], shuffle=False)

print(trainDataset.get_dataset_size())
print(testDataset.get_dataset_size())

batch_size = trainDataset.get_dataset_size()
trainDataset = trainDataset.batch(batch_size)
testDataset = testDataset.batch(batch_size)


if __name__ == "__main__":
    print(trainDataset)
    for data, label in trainDataset.create_tuple_iterator():
        print(data.dtype, data.shape)
        print(label.dtype, label.shape)
        break

    for data, label in testDataset.create_tuple_iterator():
        print(data.dtype, data.shape)
        print(label.dtype, label.shape)
        break

